/*
    pg_get_functionoid.c
        Get oids of a PostgreSQL's function by its name
    Wen Yi
*/
#include "postgres.h"
#include "fmgr.h"
#include "funcapi.h"
#include "utils/builtins.h"
#include "executor/spi.h"
#include "utils/palloc.h"
#include "utils/elog.h"
#include <string.h>

#define HELP_MESSAGE "Usage: select pg_get_functionoid('function_name');"

PG_MODULE_MAGIC;

PG_FUNCTION_INFO_V1(pg_get_functionoid);

typedef struct OidsInfo OidsInfo;
struct OidsInfo
{
    Oid *oids;
    int cursor;
    int count_of_oids;
};

Datum pg_get_functionoid(PG_FUNCTION_ARGS)
{
    FuncCallContext *function_context = NULL;
    OidsInfo *oids_info = NULL;
    if (SRF_IS_FIRSTCALL())
    {
        text *function_name = NULL;
        StringInfoData command;
        MemoryContext old_context = NULL;
        /*
            For select pg_get_functionoid();
            Return the help message
        */
        if (PG_NARGS() == 0)
        {
            text *help_message;
            help_message = (text *)palloc(VARHDRSZ + sizeof(HELP_MESSAGE));
            SET_VARSIZE(help_message, VARHDRSZ + sizeof(HELP_MESSAGE));
            memmove(VARDATA(help_message), HELP_MESSAGE, sizeof(HELP_MESSAGE));
            PG_RETURN_TEXT_P(help_message);
        }
        if (SPI_OK_CONNECT != SPI_connect())
        {
            elog(ERROR, "SPI connect failed.");
            PG_RETURN_NULL();
        }
        function_name = PG_GETARG_TEXT_P(0);
        initStringInfo(&command);
        appendStringInfo(&command, "select oid from pg_proc where proname = '%s';", text_to_cstring(function_name));
        if (SPI_execute(command.data, true, 0) != SPI_OK_SELECT)
        {
            elog(ERROR, "SPI select failed.");
            SPI_finish();
            PG_RETURN_NULL();
        }
        if (SPI_tuptable->numvals == 0)
        elog(INFO, "Complete at %d", __LINE__);
    elog(INFO, "%d %d\n", oids_info->cursor, oids_info->count_of_oids){
            SPI_finish();
            PG_RETURN_NULL();
        }
        function_context = SRF_FIRSTCALL_INIT();
        MemoryContextSwitchTo(function_context->multi_call_memory_ctx);
        oids_info = palloc(sizeof(OidsInfo));
        if (oids_info == NULL)
        {
            elog(ERROR, "Memory alloc failed.");
            SPI_finish();
            SRF_RETURN_DONE(function_context);
        }
        oids_info->oids = palloc0_array(Oid, SPI_tuptable->numvals);
        if (oids_info->oids == NULL)
        {
            elog(ERROR, "Memory alloc failed.");
            SPI_finish();
            SRF_RETURN_DONE(function_context);
        }
        oids_info->count_of_oids = SPI_tuptable->numvals;
        oids_info->cursor = 0;
        for (int i = 0; i < SPI_tuptable->numvals; i++)
        {
            oids_info->oids[i] = atoi(SPI_getvalue(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 1));
        }
        SPI_finish();
        function_context->user_fctx = (void*)oids_info;
        MemoryContextSwitchTo(old_context);
    }
    function_context = SRF_PERCALL_SETUP();
    oids_info = function_context->user_fctx;
    if (oids_info->cursor >= oids_info->count_of_oids)
        SRF_RETURN_DONE(function_context);
    else
    {
        Datum result = ObjectIdGetDatum(oids_info->oids[oids_info->cursor]);
        oids_info->cursor += 1;
        SRF_RETURN_NEXT(function_context, result);
    } 
    PG_RETURN_NULL();
}

